# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
CenterNet for training and evaluation
"""

import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import dtype as mstype
from mindspore.common.tensor import Tensor
from mindspore.common.initializer import Constant
from src.utils import Sigmoid
from src.utils import FocalLoss, RegLoss
from src.decode import DetectionDecode
from src.resnet101 import Bottleneck, ResNet101, weights_init
from .model_utils.config import dataset_config as data_cfg

BN_MOMENTUM = 0.9


def _generate_feature(cin, cout, kernel_size, head_name, head_conv=0):
    """
    Generate ResNet feature extraction function of each target head
    """
    fc = None
    if 'hm' in head_name:
        conv2d = nn.Conv2d(head_conv, cout, kernel_size=kernel_size, has_bias=True, bias_init=Constant(-2.19))
    else:
        conv2d = nn.Conv2d(head_conv, cout, kernel_size=kernel_size, has_bias=True)
    fc = nn.SequentialCell([nn.Conv2d(cin, head_conv, kernel_size=3, has_bias=True), nn.ReLU(), conv2d])
    return fc


class GatherDetectionFeatureCell(nn.Cell):
    """
    Gather ResNet features of multi-pose estimation.

    Args:
        net_config: The config info of CenterNet network.

    Returns:
        Tuple of Tensors, the target head of multi-person pose.
    """
    def __init__(self, net_config):
        super(GatherDetectionFeatureCell, self).__init__()
        self.block_class = Bottleneck
        self.layers = net_config.block_class
        heads = {'hm': data_cfg.num_classes, 'wh': 2}
        if net_config.reg_offset:
            heads.update({'reg': 2})
        head_conv = net_config.head_conv
        self.resnet101 = ResNet101(self.block_class, self.layers, heads, head_conv)

        weights_init(self.resnet101)
        self.hm_fn = _generate_feature(cin=64, cout=heads['hm'], kernel_size=1,
                                       head_name='hm', head_conv=head_conv)
        self.wh_fn = _generate_feature(cin=64, cout=heads['wh'], kernel_size=1,
                                       head_name='wh', head_conv=head_conv)
        if net_config.reg_offset:
            self.reg_fn = _generate_feature(cin=64, cout=heads['reg'], kernel_size=1,
                                            head_name='reg', head_conv=head_conv)
        self.reg_offset = net_config.reg_offset
        self.not_enable_mse_loss = not net_config.mse_loss

    def construct(self, image):
        """Defines the computation performed."""
        output = self.resnet101(image)
        feature = ()
        out = {}

        out['hm'] = self.hm_fn(output)

        out['wh'] = self.wh_fn(output)

        if self.reg_offset:
            out['reg'] = self.reg_fn(output)

        feature += (out,)
        return feature


class CenterNetLossCell(nn.Cell):
    """
    Provide object detection network losses.

    Args:
        net_config: The config info of CenterNet network.

    Returns:
        Tensor, total loss.
    """
    def __init__(self, net_config):
        super(CenterNetLossCell, self).__init__()
        self.network = GatherDetectionFeatureCell(net_config)
        self.num_stacks = net_config.num_stacks
        self.reduce_sum = ops.ReduceSum()
        self.Sigmoid = Sigmoid()
        self.FocalLoss = FocalLoss()
        self.crit = nn.MSELoss() if net_config.mse_loss else self.FocalLoss
        self.crit_reg = RegLoss(net_config.reg_loss)
        self.crit_wh = RegLoss(net_config.reg_loss)
        self.num_stacks = net_config.num_stacks
        self.wh_weight = net_config.wh_weight
        self.hm_weight = net_config.hm_weight
        self.off_weight = net_config.off_weight
        self.reg_offset = net_config.reg_offset
        self.not_enable_mse_loss = not net_config.mse_loss

    def construct(self, image, hm, reg_mask, ind, wh, reg):
        """Defines the computation performed."""
        hm_loss, wh_loss, off_loss = 0, 0, 0
        feature = self.network(image)

        for s in range(self.num_stacks):
            output = feature[s]
            if self.not_enable_mse_loss:
                output_hm = self.Sigmoid(output['hm'])
            else:
                output_hm = output['hm']
            hm_loss += self.crit(output_hm, hm) / self.num_stacks

            output_wh = output['wh']
            wh_loss += self.crit_reg(output_wh, reg_mask, ind, wh) / self.num_stacks

            if self.reg_offset and self.off_weight > 0:
                output_reg = output['reg']
                off_loss += self.crit_reg(output_reg, reg_mask, ind, reg) / self.num_stacks
        total_loss = (self.hm_weight * hm_loss + self.wh_weight * wh_loss + self.off_weight * off_loss)
        return total_loss


class ImagePreProcess(nn.Cell):
    """
    Preprocess of image on device inplace of on host to improve performance.

    Args: None

    Returns:
        Tensor, normlized images and the format were converted to be NCHW
    """
    def __init__(self):
        super(ImagePreProcess, self).__init__()
        self.transpose = ops.Transpose()
        self.perm_list = (0, 3, 1, 2)
        self.mean = Tensor(data_cfg.mean.reshape((1, 1, 1, 3)))
        self.std = Tensor(data_cfg.std.reshape((1, 1, 1, 3)))
        self.cast = ops.Cast()

    def construct(self, image):
        image = self.cast(image, mstype.float32)
        image = (image - self.mean) / self.std
        image = self.transpose(image, self.perm_list)
        return image


class CenterNetWithoutLossScaleCell(nn.Cell):
    """
    Encapsulation class of centernet training.

    Append an optimizer to the training network after that the construct
    function can be called to create the backward graph.

    Args:
        network (Cell): The training network. Note that loss function should have been added.
        optimizer (Optimizer): Optimizer for updating the weights.

    Returns:
        Tuple of Tensors, the loss, overflow flag and scaling sens of the network.
    """
    def __init__(self, network, optimizer):
        super(CenterNetWithoutLossScaleCell, self).__init__(auto_prefix=False)
        self.image = ImagePreProcess()
        self.network = network
        self.network.set_grad()
        self.weights = optimizer.parameters
        self.optimizer = optimizer
        self.grad = ops.GradOperation(get_by_list=True, sens_param=False)

    @ops.add_flags(has_effect=True)
    def construct(self, image, hm, reg_mask, ind, wh, reg):
        """Defines the computation performed."""
        image = self.image(image)
        weights = self.weights
        loss = self.network(image, hm, reg_mask, ind, wh, reg)
        grads = self.grad(self.network, weights)(image, hm, reg_mask, ind, wh, reg)
        succ = self.optimizer(grads)
        ret = loss
        return ops.depend(ret, succ)


class CenterNetWithLossScaleCell(nn.Cell):
    """
    Encapsulation class of centernet training.

    Append an optimizer to the training network after that the construct
    function can be called to create the backward graph.

    Args:
        network (Cell): The training network. Note that loss function should have been added.
        optimizer (Optimizer): Optimizer for updating the weights.
        sens (number): Static loss scale. Default: 1.

    Returns:
        Tuple of Tensors, the loss, overflow flag and scaling sens of the network.
    """
    def __init__(self, network, optimizer, sens=1):
        super(CenterNetWithLossScaleCell, self).__init__()
        self.image = ImagePreProcess()
        manager = nn.FixedLossScaleUpdateCell(loss_scale_value=sens)
        self.train_on_step = nn.TrainOneStepWithLossScaleCell(network, optimizer, scale_sense=manager)

    @ops.add_flags(has_effect=True)
    def construct(self, image, hm, reg_mask, ind, wh, reg):
        """Defines the computation performed."""
        image = self.image(image)
        return self.train_on_step(image, hm, reg_mask, ind, wh, reg)


class CenterNetDetEval(nn.Cell):
    """
    Encapsulation class of centernet testing.

    Args:
        net_config: The config info of CenterNet network.
        K(number): Max number of output objects. Default: 100.
        enable_nms_fp16(bool): Use float16 data for max_pool, adaption for CPU. Default: False.

    Returns:
        Tensor, detection of images(bboxes, score, keypoints and category id of each objects)
    """
    def __init__(self, net_config, K=100, enable_nms_fp16=False):
        super(CenterNetDetEval, self).__init__()
        self.network = GatherDetectionFeatureCell(net_config)
        self.decode = DetectionDecode(net_config, K, enable_nms_fp16)
        self.shape = ops.Shape()
        self.reshape = ops.Reshape()

    def construct(self, image):
        """Calculate prediction scores"""
        output = self.network(image)
        features = output[-1]
        detections = self.decode(features)
        return detections
